package com.woniu.utils;

import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @author 饶志
 * @version 1.0
 * @since 2025-06-11  16:52
 */
@Component
public class ConversationContextManager {

    private final Map<String, List<String>> conversationContexts = new HashMap<>();

    private static final String USER_PREFIX = "USER: ";
    private static final String AI_PREFIX = "ASSISTANT: ";

    /**
     * 添加用户的提问到上下文记录中
     * @param sessionId 用户会话ID
     * @param userInput 用户的提问
     */
    public void addUserInput(String sessionId, String userInput) {
        conversationContexts.computeIfAbsent(sessionId, k -> new ArrayList<>()).add(USER_PREFIX + userInput);
    }

    /**
     * 添加AI的回复到上下文记录中
     * @param sessionId 用户会话ID
     * @param aiResponse AI的回复
     */
    public void addAIResponse(String sessionId, String aiResponse) {
        // 合并流式响应中的换行
        String mergedResponse = aiResponse.replaceAll("\n", " ");
        conversationContexts.computeIfAbsent(sessionId, k -> new ArrayList<>()).add(AI_PREFIX + mergedResponse);
    }

    /**
     * 获取用户的上下文记录
     * @param sessionId 用户会话ID
     * @return 上下文记录
     */
    public List<String> getContext(String sessionId) {
        return conversationContexts.getOrDefault(sessionId, new ArrayList<>());
    }

    /**
     * 清除用户的上下文记录
     * @param sessionId 用户会话ID
     */
    public void clearContext(String sessionId) {
        conversationContexts.remove(sessionId);
    }

    /**
     * 添加系统消息到对话上下文
     * @param sessionId 会话ID
     * @param message 系统消息内容
     */
    public void addSystemMessage(String sessionId, String message) {
        // 实现方式取决于您现有的上下文管理逻辑
        // 以下是示例实现：

        // 获取当前会话的上下文
        List<String> context = getContext(sessionId);

        // 添加系统消息标记（可以根据需要调整格式）
        context.add("[系统] " + message);

        // 如果需要限制上下文长度
        if (context.size() > 20) { // 假设最大保留20条消息
            context.remove(0); // 移除最旧的消息
        }
    }

}
